import json
import random

import numpy as np

def load_data(mode='train'):
    # 加载JSON数据文件
    with open('2-cnn/data/mnist.json') as f:
        data = json.load(f)
    
    # 将读取到的数据区分为训练集、验证集、测试集
    train_set, val_set, test_set = data
    if mode == 'train':
        # 获得训练集
        imgs, labels = train_set[0], train_set[1]
    elif mode == 'val':
        # 获得验证集
        imgs, labels = val_set[0], val_set[1]
    elif mode == 'test':
        # 获得测试集
        imgs, labels = test_set[0], test_set[1]
    else:
        raise Exception('mode can only be one of [train, val, test]')
    print("训练数据集数量:",len(imgs))

    # 获取数据集长度
    imgs_length = len(imgs)

    # 定义数据集中每个数据的序号，根据序号读取数据
    index_list = list(range(imgs_length))
    # 读入数据时用到的批次大小
    BATCHSIZE = 100

    # 定义数据生成器
    def data_generator():
        if mode == 'train':
            # 随机打乱数据集中的数据
            random.shuffle(index_list)
        imgs_list = []
        labels_list = []

        # 遍历数据集中的每个数据
        for i in index_list:
            # 将数据转换成对应格式
            img = np.reshape(imgs[i],[1,28,28]).astype('float32')
            label = np.reshape(labels[i], [1]).astype('int64')
            imgs_list.append(img)
            labels_list.append(label)
            if len(imgs_list) == BATCHSIZE:
                # 获得一个BATCHSIZE的数据, 并返回
                yield np.array(imgs_list), np.array(labels_list)
                # 清空数据读取列表
                imgs_list = []
                labels_list = []
        # 如果剩余数据的个数小于BATCHSIZE
        # 则剩余数据一起构成一个大小为len(imgs_list)的小批次数据集mini-batch
        if len(imgs_list) > 0:
            yield np.array(imgs_list), np.array(labels_list)

    return data_generator
